"""
BEL 工具函数模块

提供数据解析、格式化、缓存、日志等通用工具函数
严格按照设计文档要求，复用现有代码
"""

import json
import re
import os
import sys
import yaml
from typing import Dict, List, Optional, TypedDict, Any, Union
from tqdm import tqdm
from loguru import logger

# 复用现有的推理模块
from inference.vllm_infer import extract_answer_try_all_methods, template, prompt, prepare_data
from vllm import LLM, SamplingParams


def prepare_data_custom(dataset: str, data_path: str, image_root: str, max_samples: int = -1):
    """
    自定义数据准备函数，处理用户数据格式
    
    Args:
        dataset: 数据集名称
        data_path: 数据文件路径
        image_root: 图片根目录
        max_samples: 最大加载样本数，-1 表示全部加载
        
    Returns:
        Tuple[List[Dict], List[Dict]]: (input_data, origin_data)
    """
    logger.info(f"使用自定义数据准备函数加载: {data_path}")
    if max_samples > 0:
        logger.info(f"限制加载样本数: {max_samples}")
    
    # 检查文件是否存在
    if not os.path.exists(data_path):
        raise FileNotFoundError(f"数据文件不存在: {data_path}")
    
    # 加载数据（按需限制数量）
    data = []
    valid_count = 0
    processed_count = 0
    
    try:
        if data_path.endswith('.json'):
            with open(data_path, 'r', encoding='utf-8') as f:
                all_data = json.load(f)
                logger.info(f"JSON 文件总计 {len(all_data)} 条数据")
                
                for item in all_data:
                    processed_count += 1
                    # 检查是否有有效答案
                    if has_valid_answer(item):
                        data.append(item)
                        valid_count += 1
                        if max_samples > 0 and valid_count >= max_samples:
                            break
                    
                    # 每处理1000条显示一次进度
                    if processed_count % 1000 == 0:
                        logger.info(f"已处理 {processed_count} 条，找到有效数据 {valid_count} 条")
                        
        elif data_path.endswith('.jsonl'):
            with open(data_path, 'r', encoding='utf-8') as f:
                for line_num, line in enumerate(f, 1):
                    try:
                        item = json.loads(line.strip())
                        processed_count += 1
                        
                        # 检查是否有有效答案
                        if has_valid_answer(item):
                            data.append(item)
                            valid_count += 1
                            if max_samples > 0 and valid_count >= max_samples:
                                logger.info(f"已达到目标样本数 {max_samples}，停止加载")
                                break
                        
                        # 每处理1000条显示一次进度
                        if processed_count % 1000 == 0:
                            logger.info(f"已处理 {processed_count} 条，找到有效数据 {valid_count} 条")
                            
                    except json.JSONDecodeError as e:
                        logger.warning(f"第 {line_num} 行 JSON 解析失败: {e}")
                        continue
                        
            logger.info(f"成功加载 JSONL 文件，处理 {processed_count} 条，有效数据 {valid_count} 条")
        else:
            raise ValueError(f"不支持的文件格式: {data_path}")
    except Exception as e:
        logger.error(f"数据文件加载失败: {e}")
        raise
    
    if not data:
        raise ValueError("没有找到有效的数据")
    
    # 打印第一条数据的结构，用于调试
    logger.info(f"第一条数据结构: {list(data[0].keys())}")
    logger.debug(f"第一条数据示例: {data[0]}")
    
    # 简单的数据格式转换，支持用户自定义格式
    input_data = []
    for i, item in enumerate(data):
        try:
            # 处理用户的数据格式：包含 image_url, instruction, output
            if 'instruction' in item:
                # 检查是否有图片，如果有图片但 instruction 中没有 <|image|> 标记，则需要添加
                has_image = ('image_url' in item and item['image_url']) or \
                           ('image' in item and item['image']) or \
                           ('image_path' in item and item['image_path'])
                
                if has_image and '<|image|>' not in item['instruction']:
                    # 使用标准的多模态模板格式
                    if template:
                        prompt_text = template.format(item['instruction'])
                        logger.debug(f"为有图片的数据添加 <|image|> 标记: {item['instruction'][:100]}...")
                    else:
                        prompt_text = item['instruction']
                else:
                    # 直接使用 instruction 作为 prompt
                    prompt_text = item['instruction']
                    logger.debug(f"使用 instruction 字段作为 prompt: {prompt_text[:100]}...")
            elif 'question' in item:
                # 使用默认格式
                question = item.get('question', '')
                if template and prompt:
                    prompt_text = template.format(prompt + question)
                else:
                    prompt_text = question
                logger.debug(f"使用 question 字段构建 prompt: {question[:100]}...")
            else:
                # 如果没有 instruction 或 question，尝试使用其他字段
                logger.warning(f"第 {i+1} 条数据缺少 instruction 或 question 字段，尝试使用其他字段")
                # 尝试使用可能的字段名
                possible_fields = ['text', 'prompt', 'content', 'problem']
                prompt_text = None
                for field in possible_fields:
                    if field in item and item[field]:
                        prompt_text = item[field]
                        logger.info(f"使用 {field} 字段作为 prompt")
                        break
                
                if prompt_text is None:
                    logger.error(f"第 {i+1} 条数据无法找到有效的文本字段")
                    logger.error(f"数据字段: {list(item.keys())}")
                    continue
            
            # 从 output 中提取 gold_answer（如果存在）
            if 'output' in item and '†Answer:' in item['output']:
                # 提取 †Answer: 后面的答案
                answer_part = item['output'].split('†Answer:')[-1].strip()
                item['gold_answer'] = answer_part
                logger.debug(f"从 output 中提取答案: {answer_part}")
            elif 'answer' in item:
                item['gold_answer'] = item['answer']
                logger.debug(f"使用 answer 字段: {item['answer']}")
            elif 'gold_answer' in item:
                # 已经有 gold_answer 字段
                logger.debug(f"使用现有的 gold_answer 字段: {item['gold_answer']}")
            else:
                logger.warning(f"第 {i+1} 条数据缺少答案字段")
            
            # 处理图片路径
            image_path = None
            if 'image_url' in item and item['image_url']:
                image_path = os.path.join(image_root, item['image_url'])
                logger.debug(f"使用 image_url 字段: {item['image_url']}")
            elif 'image' in item and item['image']:
                image_path = os.path.join(image_root, item['image'])
                logger.debug(f"使用 image 字段: {item['image']}")
            elif 'image_path' in item and item['image_path']:
                image_path = os.path.join(image_root, item['image_path'])
                logger.debug(f"使用 image_path 字段: {item['image_path']}")
            
            input_item = {'prompt': prompt_text}
            
            # 如果有图片，添加多模态数据
            if image_path and os.path.exists(image_path):
                try:
                    from PIL import Image  # type: ignore
                    input_item['multi_modal_data'] = {
                        'image': Image.open(image_path).convert('RGB')
                    }
                    logger.debug(f"成功加载图片: {image_path}")
                except Exception as e:
                    logger.warning(f"图片加载失败 {image_path}: {e}")
                    # 图片加载失败时继续使用纯文本模式
            elif image_path:
                logger.warning(f"图片文件不存在: {image_path}")
            else:
                logger.debug(f"第 {i+1} 条数据没有图片")
            
            input_data.append(input_item)
            
        except Exception as e:
            logger.error(f"处理第 {i+1} 条数据时出错: {e}")
            logger.error(f"数据内容: {item}")
            raise
    
    logger.info(f"数据准备完成，共 {len(input_data)} 条有效数据")
    return input_data, data


def has_valid_answer(item: Dict) -> bool:
    """
    检查数据项是否有有效答案
    
    Args:
        item: 数据项
        
    Returns:
        bool: 是否有有效答案
    """
    # 检查 output 字段中是否包含答案
    if 'output' in item and item['output']:
        if '†Answer:' in item['output']:
            return True
    
    # 检查是否有 answer 或 gold_answer 字段
    if 'answer' in item and item['answer']:
        return True
    if 'gold_answer' in item and item['gold_answer']:
        return True
    
    return False


class Trace(TypedDict):
    """推理轨迹数据结构"""
    question: str           # 问题文本
    steps: List[str]        # 推理步骤列表
    gold_answer: str        # 标准答案
    image: Optional[str]    # 图片路径（统一字段名）


class ProgressBar:
    """进度条封装类，支持多GPU环境"""
    
    def __init__(self, total: int, desc: str = "Processing", disable: bool = False):
        """
        初始化进度条
        
        Args:
            total: 总数
            desc: 描述
            disable: 是否禁用
        """
        self.total = total
        self.desc = desc
        self.disable = disable
        self.pbar = None
        
    def __enter__(self):
        if not self.disable:
            self.pbar = tqdm(total=self.total, desc=self.desc)
        return self
        
    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.pbar:
            self.pbar.close()
            
    def update(self, n: int = 1):
        """更新进度"""
        if self.pbar:
            self.pbar.update(n)
            
    def set_description(self, desc: str):
        """设置描述"""
        if self.pbar:
            self.pbar.set_description(desc)


def setup_logging(level: str = "INFO"):
    """
    设置日志配置
    
    Args:
        level: 日志级别
    """
    # 使用 loguru 替代标准 logging
    logger.remove()  # 移除默认处理器
    logger.add(
        sys.stdout,
        level=level,
        format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>"
    )


def load_config(config_path: str) -> Dict[str, Any]:
    """
    加载配置文件
    
    Args:
        config_path: 配置文件路径
        
    Returns:
        Dict: 配置字典
    """
    with open(config_path, 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)
    return config


def validate_config(config: Dict[str, Any]) -> bool:
    """
    验证配置文件完整性（简化版，移除自适应采样相关验证）
    
    Args:
        config: 配置字典
        
    Returns:
        bool: 是否有效
    """
    required_sections = ['model', 'data', 'algorithm', 'output']
    
    for section in required_sections:
        if section not in config:
            logger.error(f"缺少必需的配置节: {section}")
            return False
    
    # 验证模型配置
    model_config = config['model']
    if 'path' not in model_config:
        logger.error("缺少模型路径配置")
        return False
    
    # 验证模型参数范围
    if 'temperature' in model_config:
        temp = model_config['temperature']
        if not (0.0 <= temp <= 2.0):
            logger.error(f"temperature 参数超出范围 [0.0, 2.0]: {temp}")
            return False
    
    if 'top_p' in model_config:
        top_p = model_config['top_p']
        if not (0.0 <= top_p <= 1.0):
            logger.error(f"top_p 参数超出范围 [0.0, 1.0]: {top_p}")
            return False
    
    if 'max_tokens' in model_config:
        max_tokens = model_config['max_tokens']
        if max_tokens <= 0:
            logger.error(f"max_tokens 必须大于 0: {max_tokens}")
            return False
    
    if 'tensor_parallel_size' in model_config:
        tps = model_config['tensor_parallel_size']
        if tps <= 0:
            logger.error(f"tensor_parallel_size 必须大于 0: {tps}")
            return False
        
    # 验证数据配置
    data_config = config['data']
    if 'data_path' not in data_config:
        logger.error("缺少数据路径配置")
        return False
    
    # 检查数据文件是否存在
    data_path = data_config['data_path']
    if not os.path.exists(data_path):
        logger.error(f"数据文件不存在: {data_path}")
        return False
        
    # 验证算法配置（简化版）
    algorithm_config = config['algorithm']
    required_algorithm_keys = ['num_rollouts', 'n_mid', 'success_threshold']
    for key in required_algorithm_keys:
        if key not in algorithm_config:
            logger.error(f"缺少必需的算法配置: {key}")
            return False
    
    # 验证算法参数范围
    num_rollouts = algorithm_config['num_rollouts']
    n_mid = algorithm_config['n_mid']
    success_threshold = algorithm_config['success_threshold']
    
    if num_rollouts <= 0:
        logger.error(f"num_rollouts 必须大于 0: {num_rollouts}")
        return False
    
    if n_mid <= 0:
        logger.error(f"n_mid 必须大于 0: {n_mid}")
        return False
    
    if n_mid > num_rollouts:
        logger.error(f"n_mid ({n_mid}) 不能大于 num_rollouts ({num_rollouts})")
        return False
    
    if not (0.0 <= success_threshold <= 1.0):
        logger.error(f"success_threshold 参数超出范围 [0.0, 1.0]: {success_threshold}")
        return False
    
    # 验证批处理大小
    batch_size = algorithm_config.get('batch_size', 8)
    if batch_size <= 0:
        logger.error(f"batch_size 必须大于 0: {batch_size}")
        return False
    
    # 验证输出配置
    output_config = config['output']
    required_output_keys = ['rollouts', 'labeled']
    for key in required_output_keys:
        if key not in output_config:
            logger.error(f"缺少必需的输出配置: {key}")
            return False
    
    # 验证性能配置（可选）
    if 'performance' in config:
        performance_config = config['performance']
        
        if 'gpu_memory_utilization' in performance_config:
            gpu_mem = performance_config['gpu_memory_utilization']
            if not (0.0 < gpu_mem <= 1.0):
                logger.error(f"gpu_memory_utilization 参数超出范围 (0.0, 1.0]: {gpu_mem}")
                return False
        
        if 'max_model_len' in performance_config:
            max_len = performance_config['max_model_len']
            if max_len <= 0:
                logger.error(f"max_model_len 必须大于 0: {max_len}")
                return False
    
    logger.info("配置文件验证通过")
    return True


def create_llm_instance(config: Dict[str, Any]) -> LLM:
    """
    统一的LLM初始化函数，完全复用 inference/vllm_infer.py 的成功做法
    
    Args:
        config: 配置字典
        
    Returns:
        LLM: 初始化的LLM实例
    """
    model_config = config['model']
    model_path = model_config['path']
    
    logger.info("🚀 开始初始化 LLM 实例...")
    logger.info(f"📁 模型路径: {model_path}")
    
    # 设置多进程启动方式为 spawn（解决CUDA多进程问题）
    tensor_parallel_size = model_config.get('tensor_parallel_size', 1)
    if tensor_parallel_size > 1:
        os.environ['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn'
        logger.info("🔧 设置多进程启动方式为 spawn（多GPU支持）")
    
    # 设置环境变量（如果配置中有指定）
    env_config = config.get('environment', {})
    if ('cuda_visible_devices' in env_config
            and model_config.get('distributed_executor_backend', None) != 'ray'):
        cuda_devices = env_config['cuda_visible_devices']
        os.environ['CUDA_VISIBLE_DEVICES'] = cuda_devices
        logger.info(f"🔧 设置 CUDA_VISIBLE_DEVICES={cuda_devices}")
    
    # 注册 URSA 模型（使用 __init__.py 中的注册逻辑）
    logger.info("⏳ 正在加载模型，请稍等...")
    
    try:
        # 获取性能配置
        performance_config = config.get('performance', {})
        
        # GPU 可用性检测与并行度自适应
        distributed_backend = model_config.get('distributed_executor_backend', None)
        if distributed_backend == 'ray':
            try:
                import torch  # type: ignore
                available_gpus: int = torch.cuda.device_count()
            except Exception:
                available_gpus = 0

            if available_gpus == 0:
                logger.warning("⚠️ 未检测到可用的 CUDA 设备，后续初始化可能失败")
            elif tensor_parallel_size > available_gpus:
                logger.warning(
                    "⚠️ 检测到张量并行大小 (%d) 大于实际可用 GPU 数 (%d)，自动降级以避免 'device < num_gpus' 错误",
                    tensor_parallel_size, available_gpus)
                tensor_parallel_size = available_gpus
            
            cuda_visible = os.environ.get('CUDA_VISIBLE_DEVICES')
            if cuda_visible:
                visible_list = [d for d in cuda_visible.split(',') if d.strip()]
                if len(visible_list) < tensor_parallel_size:
                    logger.warning(
                        "⚠️ 环境变量 CUDA_VISIBLE_DEVICES 提供的 GPU 数量 (%d) 小于张量并行大小 (%d)，自动按可见设备数量调整",
                        len(visible_list), tensor_parallel_size)
                    tensor_parallel_size = len(visible_list)
        
        # 创建 LLM 实例
        llm = LLM(
            model=model_path,
            tensor_parallel_size=tensor_parallel_size,
            gpu_memory_utilization=performance_config.get('gpu_memory_utilization', 0.85),
            max_model_len=performance_config.get('max_model_len', 4096),
            enable_prefix_caching=performance_config.get('enable_prefix_caching', True),
            enforce_eager=performance_config.get('enforce_eager', False),
            dtype=model_config.get('dtype', 'auto'),
            quantization=model_config.get('quantization', None),
            disable_custom_all_reduce=model_config.get('disable_custom_all_reduce', False),
            distributed_executor_backend=model_config.get('distributed_executor_backend', None)
        )
        logger.info("✅ LLM 初始化完成!")
        logger.info(f"📊 使用 {tensor_parallel_size} 个GPU进行张量并行")
        logger.info(f"🔧 GPU内存使用率: {performance_config.get('gpu_memory_utilization', 0.85)}")
        logger.info(f"🚫 禁用自定义AllReduce: {model_config.get('disable_custom_all_reduce', False)}")
        logger.info(f"🌐 分布式后端: {model_config.get('distributed_executor_backend', 'mp')}")
        return llm
    except Exception as e:
        logger.error(f"❌ LLM 初始化失败: {e}")
        logger.error("请检查模型路径和配置参数")
        raise


def create_sampling_params(config: Dict[str, Any], n: Optional[int] = None) -> SamplingParams:
    """
    创建采样参数，修复参数设置问题
    
    Args:
        config: 配置字典
        n: 采样数量，如果提供则覆盖配置中的值
        
    Returns:
        SamplingParams: 采样参数
    """
    model_config = config['model']
    algorithm_config = config['algorithm']
    
    # 使用传入的 n 值或配置中的 num_rollouts
    sampling_n = n if n is not None else algorithm_config['num_rollouts']
    
    # 根据 vLLM 文档使用正确的参数名
    return SamplingParams(
        temperature=model_config.get('temperature', 0.2),
        top_p=model_config.get('top_p', 0.9),
        max_tokens=model_config.get('max_tokens', 2048),
        n=sampling_n  # 生成多个序列
    )


def validate_answer(reasoning: str, gold_answer: str) -> bool:
    """
    统一的答案验证函数，直接复用现有的答案提取逻辑
    
    Args:
        reasoning: 推理文本
        gold_answer: 标准答案
        
    Returns:
        bool: 是否正确
    """
    try:
        # 直接复用现有的答案提取函数
        extracted_answer = extract_answer_try_all_methods(reasoning)
        
        # 简单的答案比较
        extracted_clean = extracted_answer.strip().lower()
        gold_clean = gold_answer.strip().lower()
        
        return extracted_clean == gold_clean
        
    except Exception as e:
        logger.debug(f"答案验证失败: {e}")
        return False


def count_lines(file_path: str) -> int:
    """
    统计文件行数 - 用于断点续跑
    
    Args:
        file_path: 文件路径
        
    Returns:
        int: 行数
    """
    if not os.path.exists(file_path):
        return 0
    
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            return sum(1 for _ in f)
    except Exception as e:
        logger.error(f"统计文件行数失败: {e}")
        return 0


def split_steps(output: str) -> List[str]:
    """
    基于正则表达式拆分推理步骤
    
    Args:
        output: 推理输出文本
        
    Returns:
        List[str]: 步骤列表
    """
    # 使用正则表达式匹配 Step 前缀
    step_pattern = r'(?i)step\s*\d+\s*[:\.]?'
    matches = list(re.finditer(step_pattern, output))
    
    if not matches:
        # 如果没有匹配到 Step 前缀，按换行拆分
        lines = output.strip().split('\n')
        return [line.strip() for line in lines if line.strip()]
    
    steps = []
    for i, match in enumerate(matches):
        start = match.start()
        if i + 1 < len(matches):
            end = matches[i + 1].start()
            step_content = output[start:end].strip()
        else:
            step_content = output[start:].strip()
        
        if step_content:
            steps.append(step_content)
    
    return steps


def insert_labels(output: str, labels: List[int]) -> str:
    """
    在推理步骤中插入标签
    
    Args:
        output: 原始推理输出
        labels: 标签列表 (1=正确, 0=错误)
        
    Returns:
        str: 插入标签后的输出
    """
    steps = split_steps(output)
    
    if len(labels) != len(steps):
        logger.warning(f"标签数量({len(labels)})与步骤数量({len(steps)})不匹配")
        return output
    
    labeled_parts = []
    for i, (step, label) in enumerate(zip(steps, labels)):
        if label == 1:
            labeled_parts.append(f"{step} <pos>")
        else:
            labeled_parts.append(f"{step} <neg>")
    
    return '\n'.join(labeled_parts)


def extract_gold_answer(data: Dict) -> str:
    """
    提取标准答案
    
    Args:
        data: 数据字典
        
    Returns:
        str: 标准答案
    """
    # 尝试多种字段名
    answer_fields = ['gold_answer', 'answer', 'correct_answer', 'target']
    
    for field in answer_fields:
        if field in data and data[field]:
            return str(data[field]).strip()
    
    # 如果没有找到标准答案，尝试从输出中提取
    if 'output' in data:
        return extract_answer_try_all_methods(data['output'])
    
    return ""


def normalize_image_path(data: Dict) -> str:
    """
    统一图片路径字段处理
    
    Args:
        data: 数据字典
        
    Returns:
        str: 图片路径
    """
    # 统一处理不同的图片路径字段名
    image_fields = ['image', 'image_url', 'image_path', 'img_path']
    
    for field in image_fields:
        if field in data and data[field]:
            return str(data[field]).strip()
    
    return ""


def ensure_dir(path: str):
    """
    确保目录存在
    
    Args:
        path: 文件路径或目录路径
    """
    # 判断是否为文件路径（通过扩展名判断）
    if os.path.splitext(path)[1]:  # 有文件扩展名，说明是文件路径
        dir_path = os.path.dirname(path)
    else:  # 没有扩展名，说明是目录路径
        dir_path = path
    
    if dir_path and not os.path.exists(dir_path):
        os.makedirs(dir_path, exist_ok=True)


def format_steps_with_labels(steps: List[str], labels: List[int]) -> str:
    """
    格式化带标签的步骤
    
    Args:
        steps: 步骤列表
        labels: 标签列表
        
    Returns:
        str: 格式化后的字符串
    """
    if len(steps) != len(labels):
        logger.warning(f"步骤数量({len(steps)})与标签数量({len(labels)})不匹配")
        return '\n'.join(steps)
    
    formatted_steps = []
    for i, (step, label) in enumerate(zip(steps, labels)):
        prefix = f"Step {i+1}: "
        if not step.startswith(prefix):
            step = prefix + step
        
        if label == 1:
            formatted_steps.append(f"{step} <pos>")
        else:
            formatted_steps.append(f"{step} <neg>")
    
    return '\n'.join(formatted_steps)


def parse_trace(data: Dict) -> Trace:
    """
    解析推理轨迹数据
    
    Args:
        data: 原始数据字典
        
    Returns:
        Trace: 解析后的轨迹数据
    """
    # 提取问题
    question = data.get('question', data.get('instruction', ''))
    
    # 提取步骤
    output = data.get('output', '')
    steps = split_steps(output) if output else []
    
    # 提取标准答案
    gold_answer = extract_gold_answer(data)
    
    # 统一处理图片路径
    image = normalize_image_path(data)
    
    return Trace(
        question=question,
        steps=steps,
        gold_answer=gold_answer,
        image=image
    )


def validate_bel_format(data: Dict) -> bool:
    """
    验证BEL数据格式
    
    Args:
        data: 数据字典
        
    Returns:
        bool: 是否符合格式要求
    """
    required_fields = ['rollouts']
    
    for field in required_fields:
        if field not in data:
            logger.error(f"缺少必需字段: {field}")
            return False
    
    # 验证rollouts格式
    rollouts = data['rollouts']
    if not isinstance(rollouts, list):
        logger.error("rollouts 必须是列表")
        return False
    
    for i, rollout in enumerate(rollouts):
        if not isinstance(rollout, dict):
            logger.error(f"rollout {i} 必须是字典")
            return False
        
        if 'output' not in rollout:
            logger.error(f"rollout {i} 缺少 output 字段")
            return False
    
    return True 